/**
 * 
 */
package edu.umd.clip.lm.model;

import java.util.*;
import java.io.*;

import com.sleepycat.bind.tuple.*;

import edu.berkeley.nlp.util.MutableInteger;
import edu.umd.clip.lm.factors.FactorTupleDescription;
import edu.umd.clip.lm.util.tree.*;
import edu.umd.clip.lm.util.*;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class DecoderCompactProbTree implements Cloneable, Externalizable /*extends AbstractCompactProbTree*/ {
	public static final short NO_LINK = -1; 
	private short root = 0;
	private Data data; 
	private static DecoderCompactProbTree emptyTree;
 
	private static class Data implements Cloneable, Externalizable {
		public short[] leftLinks;
		public short[] rightLinks;
		public short[] parentLinks;
		public float[] nodes;
		
		private Data() {}
		
		private Data(ObjectInput in) throws IOException, ClassNotFoundException {
			readExternal(in);
		}
		
		private Data(TupleInput in) throws IndexOutOfBoundsException {
			readTuple(in);
		}
		
		private Data(int size) {
			leftLinks = new short[size];
			Arrays.fill(leftLinks, NO_LINK);
			
			rightLinks = new short[size];
			Arrays.fill(rightLinks, NO_LINK);
			
			parentLinks = new short[size];
			Arrays.fill(parentLinks, NO_LINK);
			
			nodes = new float[size];
			Arrays.fill(nodes, (float)0.0);			
		}

		/* (non-Javadoc)
		 * @see java.lang.Object#clone()
		 */
		@Override
		protected Object clone() {
			Data copy = new Data();
			copy.leftLinks = Arrays.copyOf(leftLinks, leftLinks.length);
			copy.rightLinks = Arrays.copyOf(rightLinks, rightLinks.length);
			copy.parentLinks = Arrays.copyOf(parentLinks, parentLinks.length);
			copy.nodes = Arrays.copyOf(nodes, nodes.length);
			return copy;
		}

		/* (non-Javadoc)
		 * @see java.lang.Object#hashCode()
		 */
		@Override
		public int hashCode() {
			return Arrays.hashCode(nodes);
		}

		public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
			int size = in.readInt();
			leftLinks = new short[size];
			for(int i=0; i<size; ++i) {
				leftLinks[i] = in.readShort();
			}
			
			rightLinks = new short[size];
			for(int i=0; i<size; ++i) {
				rightLinks[i] = in.readShort();
			}
			
			parentLinks = new short[size];
			for(int i=0; i<size; ++i) {
				parentLinks[i] = in.readShort();
			}
			
			nodes = new float[size];
			for(int i=0; i<size; ++i) {
				nodes[i] = in.readFloat();
			}
		}

		public void writeExternal(ObjectOutput out) throws IOException {
			out.writeInt(nodes.length);
			for(short v : leftLinks) {
				out.writeShort(v);
			}
			
			for(short v : rightLinks) {
				out.writeShort(v);
			}
			
			for(short v : parentLinks) {
				out.writeShort(v);
			}
			
			for(float v : nodes) {
				out.writeFloat(v);
			}
		}

		public void readTuple(TupleInput in) throws IndexOutOfBoundsException {
			int size = in.readInt();
			leftLinks = new short[size];
			for(int i=0; i<size; ++i) {
				leftLinks[i] = in.readShort();
			}
			
			rightLinks = new short[size];
			for(int i=0; i<size; ++i) {
				rightLinks[i] = in.readShort();
			}
			
			parentLinks = new short[size];
			for(int i=0; i<size; ++i) {
				parentLinks[i] = in.readShort();
			}
			
			nodes = new float[size];
			for(int i=0; i<size; ++i) {
				nodes[i] = in.readFloat();
			}
		}

		public void writeTuple(TupleOutput out) {
			out.writeInt(nodes.length);
			for(short v : leftLinks) {
				out.writeShort(v);
			}
			
			for(short v : rightLinks) {
				out.writeShort(v);
			}
			
			for(short v : parentLinks) {
				out.writeShort(v);
			}
			
			for(float v : nodes) {
				out.writeFloat(v);
			}
		}
	}
	
	/**
	 * @param root
	 * @param data
	 */
	private DecoderCompactProbTree(short root, Data data) {
		this.root = root;
		this.data = data;
	}

	protected DecoderCompactProbTree() {
		data = new Data();
	}
	
	protected DecoderCompactProbTree(int size) {
		data = new Data(size);
	}
	
	public DecoderCompactProbTree(ObjectInput in) throws IOException, ClassNotFoundException {
		readExternal(in);
	}
	
	/**
	 * @param in
	 */
	public DecoderCompactProbTree(TupleInput in) throws IndexOutOfBoundsException {
		readTuple(in);
	}

	protected static DecoderCompactProbTree getSingleNodeTree() {
		DecoderCompactProbTree tree = new DecoderCompactProbTree(1);
		tree.data.nodes[0] = (float)1.0;
		return tree;
	}
	
	public static <T extends XMLizable> DecoderCompactProbTree constructEmptyTree(BinaryTree<T> bTree) {
		final HashMap<BinaryTree<T>, Integer> nodes = 
			new HashMap<BinaryTree<T>, Integer>(100);
		{
			BinaryTreeFinder<T> finder = new BinaryTreeFinder<T>() {
				public boolean apply(BinaryTree<T> node) {
					nodes.put(node, new Integer(nodes.size()));
					return false;
				}
			};
			bTree.searchPreorder(finder);
		}
		assert(nodes.size() < Short.MAX_VALUE);
		
		final DecoderCompactProbTree tree = new DecoderCompactProbTree(nodes.size());
		
		{
			BinaryTreeFinder<T> finder = new BinaryTreeFinder<T>() {
				public boolean apply(BinaryTree<T> node) {
					short nodeIdx = nodes.get(node).shortValue();
					if (node.getParent() != null) {
						short parentIdx = nodes.get(node.getParent()).shortValue();
						tree.data.parentLinks[nodeIdx] = parentIdx; 
					}
					
					if (node.getLeft() != null) {
						short leftChildIdx = nodes.get(node.getLeft()).shortValue();
						tree.data.leftLinks[nodeIdx] = leftChildIdx;
					}
					if (node.getRight() != null) {
						short rightChildIdx = nodes.get(node.getRight()).shortValue();
						tree.data.rightLinks[nodeIdx] = rightChildIdx;
					}
					return false;
				}
			};
			bTree.searchPostorder(finder);
		}
		return tree;
	}
	
	// cut branches
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#prune(double)
	 */
	public DecoderCompactProbTree prune(final double threshold) {
		
		DecoderCompactProbTree tree = (DecoderCompactProbTree)clone();
		short numLeaves = 0;
		for(short i=0; i<tree.data.nodes.length; i++) {
			// clean inner nodes
			if (data.leftLinks[i] != NO_LINK || data.rightLinks[i] != NO_LINK) {
				tree.data.nodes[i] = (float)0.0;
				continue;
			}
			// remove leaves under the threshold
			if (tree.data.nodes[i] < threshold) {
				tree.data.nodes[i] = (float)0.0;
				++numLeaves;
			}
		}
		if (numLeaves == 0) {
			return null;
		}
		
		for(short i=(short)(tree.data.nodes.length-1); i>0; --i) {
			// the only node without a parent is the root (i==0)
			short parent = getParent(i);
			tree.data.nodes[parent] += tree.data.nodes[i];
			if (tree.data.nodes[i] == 0.0) {
				tree.data.parentLinks[i] = NO_LINK;
			}
			if (tree.hasLeft(i) && tree.data.nodes[tree.data.leftLinks[i]] == 0.0) {
				tree.data.leftLinks[i] = NO_LINK;
			}
			if (tree.hasRight(i) && tree.data.nodes[tree.data.rightLinks[i]] == 0.0) {
				tree.data.rightLinks[i] = NO_LINK;
			}
		}
		return tree;
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#compact()
	 */
	public DecoderCompactProbTree compact() {
		short numNodes = 0;
		short[] nodeMap = new short[data.nodes.length];
		for(short i = 0; i<data.nodes.length; ++i) {
			if (data.nodes[i] != 0.0) {
				nodeMap[i] = numNodes++;
			} else {
				nodeMap[i] = NO_LINK;
			}
		}
		if (numNodes == data.nodes.length) return this;
		if (numNodes == 0) {
			return null;
		}
		
		DecoderCompactProbTree tree = new DecoderCompactProbTree(numNodes);
		for(short i=0; i<data.nodes.length; ++i) {
			short map = nodeMap[i];
			if (map != NO_LINK) {
				tree.data.nodes[map] = data.nodes[i];
				if (hasLeft(i)) {
					tree.data.leftLinks[map] = nodeMap[getLeft(i)];
				}
				if (hasRight(i)) {
					tree.data.rightLinks[map] = nodeMap[getRight(i)];
				}
				if (data.parentLinks[i] != NO_LINK) {
					tree.data.parentLinks[map] = nodeMap[getParent(i)];
				}
			}
		}
		return tree;
	}

	public DecoderCompactProbTree copy() {
		return copySubtree(root);
	}
	
	public DecoderCompactProbTree copySubtree(short root) {
		final short[] nodesMap = new short[data.nodes.length-root];
		Arrays.fill(nodesMap, (short)-1);
		
		final MutableInteger nodeCount = new MutableInteger(0);
		
		DecoderCompactProbTreeFinder finder = new DecoderCompactProbTreeFinder(this) {
			@Override
			public boolean apply(short node) {
				nodesMap[node - tree.root] = nodeCount.shortValue();
				nodeCount.add(1);
				return false;
			}
		};
		searchPreorder(finder, root);
		
		Data data = new Data(nodeCount.shortValue());
		Arrays.fill(data.leftLinks, NO_LINK);
		Arrays.fill(data.rightLinks, NO_LINK);
		Arrays.fill(data.parentLinks, NO_LINK);
		for(short i = 0; i < nodesMap.length; ++i) {
			short dst = nodesMap[i];
			if (dst == NO_LINK) continue;
			short src = (short) (i + root);
			data.nodes[dst] = this.data.nodes[src];
			short left = (short) (this.data.leftLinks[src]-root);
			if (left >= 0) {
				data.leftLinks[dst] = nodesMap[left];
			}
			short right = (short) (this.data.rightLinks[src]-root);
			if (right >= 0) {
				data.rightLinks[dst] = nodesMap[right];
			}
			short parent = (short) (this.data.parentLinks[src]-root);
			if (parent >= 0) {
				data.parentLinks[dst] = nodesMap[parent];
			}
		}
		DecoderCompactProbTree tree = new DecoderCompactProbTree((short)0, data);
		return tree;
	}
	
	private static void addTree(DecoderCompactProbTree t1, short pos1, DecoderCompactProbTree t2, short pos2) {
		t1.data.nodes[pos1] += t2.data.nodes[pos2];
		if (t2.hasLeft(pos2)) {
			addTree(t1, t1.getLeft(pos1), t2, t2.getLeft(pos2));
		}
		if (t2.hasRight(pos2)) {
			addTree(t1, t1.getRight(pos1), t2, t2.getRight(pos2));
		}
	}
	
	protected void addTrees(Collection<DecoderCompactProbTree> trees) {
		for(DecoderCompactProbTree t : trees) {
			addTree(this, root, t, t.root);
		}
	}
	
	protected void addTree(DecoderCompactProbTree t) {
		addTree(this, root, t, t.root);
	}
	
	private boolean equals(short pos, short otherPos, DecoderCompactProbTree other) {
		if (data.nodes[pos] != other.data.nodes[otherPos]) return false;
		
		short left = getLeft(pos);
		short otherLeft = other.getLeft(otherPos);
		if (left != otherLeft) return false;
		if (left != NO_LINK && !equals(left, otherLeft, other)) return false;
		
		short right = getRight(pos);
		short otherRight = other.getRight(otherPos);
		if (right != otherRight) return false;
		if (right != NO_LINK && !equals(right, otherRight, other)) return false;
		
		return true;
	}
	
	private void addTwoTrees(short myPos, DecoderCompactProbTree t1, short pos1, DecoderCompactProbTree t2, short pos2) {
		data.nodes[myPos] += t1.data.nodes[pos1];
		data.nodes[myPos] += t2.data.nodes[pos2];

		short left1 = t1.getLeft(pos1);
		short left2 = t2.getLeft(pos2);
		if (left1 != NO_LINK && left2 != NO_LINK) {
			addTwoTrees(getLeft(myPos), t1, left1, t2, left2);
		} else if (left1 != NO_LINK) {
			addTree(this, getLeft(myPos), t1, left1);
		} else if (left2 != NO_LINK) {
			addTree(this, getLeft(myPos), t2, left2);
		}
		
		short right1 = t1.getRight(pos1);
		short right2 = t2.getRight(pos2);
		
		if (right1 != NO_LINK && right2 != NO_LINK) {
			addTwoTrees(getRight(myPos), t1, right1, t2, right2);
		} else if (right1 != NO_LINK) {
			addTree(this, getRight(myPos), t1, right1);
		} else if (right2 != NO_LINK) {
			addTree(this, getRight(myPos), t2, right2);
		}
	}
	
	protected void addTwoTrees(DecoderCompactProbTree tree1, DecoderCompactProbTree tree2) {
		// TODO doesn't seem to improve performance much
		addTwoTrees(root, tree1, tree1.root, tree2, tree2.root);
	}
	
	@Override
	public int hashCode() {
		final int prime = 31;
		int result = 1;
		result = prime * result + Float.floatToIntBits(data.nodes[root]);
		result = prime * result + root;
		return result;
	}

	@Override
	public boolean equals(Object obj) {
		if (this == obj)
			return true;
		if (obj == null)
			return false;
		if (!(obj instanceof DecoderCompactProbTree))
			return false;
		DecoderCompactProbTree other = (DecoderCompactProbTree) obj;
		if (root != other.root)
			return false;
		return equals(root, other.root, other);
	}

	public short findNode(BinaryPrefix path) {
		short node = root;
		for(int i=0; i<path.getSize(); ++i) {
			if (path.getBits().get(i)) {
				node = getRight(node);
			} else {
				node = getLeft(node);
			}
			if (node == NO_LINK) return NO_LINK;
		}
		return node;
	}

	public int countLeaves() {
		int nrLeaves = 0;
		for(int i=0; i<data.leftLinks.length; ++i) {
			if (data.leftLinks[i] == NO_LINK && data.rightLinks[i] == NO_LINK) {
				++nrLeaves;
			}
		}
		return nrLeaves;
	}
	
	/*
	public ArrayList<PathProb> getPaths(int sizeHint) {
		ArrayList<PathProb> paths = new ArrayList<PathProb>(sizeHint);
		Experiment experiment = Experiment.getInstance();
		
		for(int i=0; i<data.leftLinks.length; ++i) {
			if (data.leftLinks[i] == NO_LINK && data.rightLinks[i] == NO_LINK) {
				int item = i;
				float prob = data.nodes[i];
				BinaryPrefix prefix = new BinaryPrefix(0);
				while(item != NO_LINK) {
					int parent = data.parentLinks[item];
					if (parent != NO_LINK) {
						prefix = prefix.append(data.rightLinks[parent] == item);
					}
					item = parent;
				}
				prefix = prefix.reverse();
				paths.add(new PathProb(experiment.getCachedPrefix(prefix), prob));
			}
		}
		return paths;
	}
	*/
	
	public BinaryPrefix getSinglePath() {
		short node = root;
		BinaryPrefix path = new BinaryPrefix(0);
		while(true) {
			if (hasLeft(node)) {
				if (hasRight(node))	return null;
				path = path.appendZero();
				node = data.leftLinks[node];
			} else if (hasRight(node)) {
				if (hasLeft(node)) return null;
				path = path.appendOne();
				node = data.rightLinks[node];
			} else {
				break;
			}
		}
		return path;
	}
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#getProbability(edu.umd.clip.lm.util.BinaryPrefix)
	 */
	public double getProbability(BinaryPrefix path) {
		short node = findNode(path);
		if (node == NO_LINK) return 0;
		return data.nodes[node];
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#setProbability(edu.umd.clip.lm.util.BinaryPrefix, double)
	 */
	public void setProbability(BinaryPrefix path, double prob) {
		short node = findNode(path);
		assert(node != NO_LINK);
		if (node == NO_LINK) return;
		double diff = prob - data.nodes[node];
		do {
			data.nodes[node] += diff;
			node = data.parentLinks[node];
		} while(node != NO_LINK);
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#addProbability(edu.umd.clip.lm.util.BinaryPrefix, double)
	 */
	public void addProbability(BinaryPrefix path, double prob) {
		short node = root;
		for(int i=0; i<path.getSize(); ++i) {
			data.nodes[node] += prob;
			if (path.getBits().get(i)) {
				node = getRight(node);
			} else {
				node = getLeft(node);
			}
			if (node == NO_LINK) {
				System.out.printf("unexpected leaf: %s, i=%d\n", path.toString(), i);
				return;
			}
		}
		data.nodes[node] += prob;
		
		if (hasLeft(node) || hasRight(node)) {
			System.out.printf("Not a leaf: %s, node=%d\n", path.toString(), node);
		}
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#checkTree()
	 */
	public boolean checkTree() {
		boolean success = true;
		/*
		for(short i=0; i<data.nodes.length; ++i) {
			double sum = 0;
			if (hasLeft(i)) {
				sum += data.nodes[getLeft(i)];
			}
			if (hasRight(i)) {
				sum += data.nodes[getRight(i)];
			}
			if (!ProbMath.approxEqual(sum, data.nodes[i])) {
				System.out.printf("prob at a node is not equivalent to the sum of its children (%d)\n", i);
				success = false;
			}
		}*/
		return success;
	}
	/**
	 * cuts off a tree at the node where path points to
	 * both resulting subtrees share data structures, thus compact(), prune(), etc., must not be called 
	 */
	protected DecoderCompactProbTree cut(BinaryPrefix path) {
		short node = findNode(path);
		if (node == NO_LINK) return null;
		float subTreeProb = data.nodes[node];

		short parent = data.parentLinks[node];

		if (parent == NO_LINK) {
			// the whole tree is cut
			return this;
		}
		
		for(short i = parent; i != NO_LINK; i = data.parentLinks[i]) {
			data.nodes[i] -= subTreeProb;
			if (data.nodes[i] < ProbMath.MIN_FLOAT_PRECISION) {
				// eliminate precision noise
				data.nodes[i] = 0;
			}
		}

		if (data.leftLinks[parent] == node) {
			data.leftLinks[parent] = NO_LINK;
		} else if (data.rightLinks[parent] == node) {
			data.rightLinks[parent] = NO_LINK;
		}
		data.parentLinks[node] = NO_LINK;
		
		if (subTreeProb < ProbMath.MIN_FLOAT_PRECISION) {
			//System.out.printf("subTreeProb=%f, parentProb=%f\n", subTreeProb, data.nodes[parent]);
			return null;
		}
		
		DecoderCompactProbTree tree = new DecoderCompactProbTree(node, data);
		return tree;
	}
	
	/* (non-Javadoc)
	 * @see java.lang.Object#clone()
	 */
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#clone()
	 */
	@Override
	public Object clone() {
		DecoderCompactProbTree copy = new DecoderCompactProbTree();
		copy.data = (Data) data.clone();
		copy.root = root;
		return copy;
	}

	/**
	 * Multiplies probabilities by scale
	 * @param scale
	 */
	public void scale(double scale) {
		for(int i=0; i<data.nodes.length; ++i) {
			data.nodes[i] *= scale;
		}
	}
	
	public short getLeft(short node) {
		return data.leftLinks[node];
	}
	
	public boolean hasLeft(short node) {
		return getLeft(node) != NO_LINK;
	}
	
	public short getRight(short node) {
		return data.rightLinks[node];
	}
	
	public boolean hasRight(short node) {
		return getRight(node) != NO_LINK;
	}
	
	public short getParent(short node) {
		return data.parentLinks[node];
	}
	
	public double getProb(short node) {
		return data.nodes[node];
	}
	
	public void setProb(short node, float prob) {
		data.nodes[node] = prob;
	}
	
	public double getTotalProb() {
		return data.nodes[root];
	}
	
	public short searchPreorder(DecoderCompactProbTreeFinder finder) {
		return searchPreorder(finder, root);
	}
	
	public short searchPostorder(DecoderCompactProbTreeFinder finder) {
		return searchPostorder(finder, root);
	}
	
	public short searchPreorder(DecoderCompactProbTreeFinder finder, short node) {
		if (finder.apply(node)) return node;
		if (hasLeft(node)) {
			short result = searchPreorder(finder, getLeft(node));
			if (result != NO_LINK) return result;
		}
		if (hasRight(node)) {
			short result = searchPreorder(finder, getRight(node));
			if (result != NO_LINK) return result;
		}
		return NO_LINK;
	}

	public short searchPostorder(DecoderCompactProbTreeFinder finder, short node) {
		if (hasLeft(node)) {
			short result = searchPostorder(finder, getLeft(node));
			if (result != NO_LINK) return result;
		}
		if (hasRight(node)) {
			short result = searchPostorder(finder, getRight(node));
			if (result != NO_LINK) return result;
		}
		if (finder.apply(node)) return node;
		return NO_LINK;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#readExternal(java.io.ObjectInput)
	 */
	public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
		root = in.readShort();
		data = new Data(in);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#writeExternal(java.io.ObjectOutput)
	 */
	public void writeExternal(ObjectOutput out) throws IOException {
		//super.writeExternal(out);
		out.writeShort(root);
		data.writeExternal(out);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.AbstractCompactProbTree#readTuple(com.sleepycat.bind.tuple.TupleInput)
	 */
	
	public void readTuple(TupleInput in) throws IndexOutOfBoundsException {
		root = in.readShort();
		data = new Data(in);
	}

	
	public void writeTuple(TupleOutput out) {
		//super.writeTuple(out);
		out.writeShort(root);
		data.writeTuple(out);
	}
	
	public DecoderCompactProbTree getCompactProbTree() {
		return this;
	}

	public static DecoderCompactProbTree getEmptyTree() {
		return emptyTree;
	}

	public static void setEmptyTree(DecoderCompactProbTree emptyTree) {
		DecoderCompactProbTree.emptyTree = emptyTree;
	}

}
